import numpy as np  # 用于数值计算和数组操作
import faiss  # 用于高效的相似性搜索和聚类
from util import createXY  # 自定义函数，用于创建特征(X)和标签(Y)
from sklearn.model_selection import train_test_split  # 用于将数据集分割为训练集和测试集
from sklearn.neighbors import KNeighborsClassifier  # 传统的K近邻分类器
import argparse  # 用于解析命令行参数
import logging  # 用于记录程序运行过程中的日志信息
from tqdm import tqdm  # 用于在控制台显示进度条
from FaissKNeighbors import FaissKNeighbors  # 自定义的基于Faiss的K近邻分类器
import os  # 用于与操作系统交互，如读取环境变量
import sys  # 用于访问Python解释器相关的变量和函数

# 设置日志记录的格式和级别
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def get_args():
    """
    解析并返回命令行参数。
    
    返回:
        argparse.Namespace: 包含解析后的命令行参数
    """
    parser = argparse.ArgumentParser(description='使用CPU或GPU训练模型。')
    parser.add_argument('-m', '--mode', type=str, required=True, choices=['cpu', 'gpu'], help='选择训练模式：CPU或GPU。')
    parser.add_argument('-f', '--feature', type=str, required=True, choices=['flat', 'vgg'], help='选择特征提取方法：flat或vgg。')
    parser.add_argument('-l', '--library', type=str, required=True, choices=['sklearn', 'faiss'], help='选择使用的库：sklearn或faiss。')
    return parser.parse_args()

def check_cuda_environment():
    """
    检查CUDA环境配置，包括环境变量和PyTorch的CUDA支持情况。
    """
    logging.info("检查 CUDA 环境...")
    # 检查CUDA_VISIBLE_DEVICES环境变量
    if 'CUDA_VISIBLE_DEVICES' in os.environ:
        logging.info(f"CUDA_VISIBLE_DEVICES: {os.environ['CUDA_VISIBLE_DEVICES']}")
    else:
        logging.warning("未设置 CUDA_VISIBLE_DEVICES 环境变量")

    # 尝试导入PyTorch并检查CUDA可用性
    try:
        import torch
        logging.info(f"PyTorch 版本: {torch.__version__}")
        logging.info(f"CUDA 是否可用: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            logging.info(f"CUDA 版本: {torch.version.cuda}")
            logging.info(f"可用 GPU 数量: {torch.cuda.device_count()}")
    except ImportError:
        logging.warning("未安装 PyTorch，无法检查 CUDA 可用性")

def check_gpu_support():
    """
    检查FAISS库的GPU支持情况。
    
    返回:
        bool: 如果支持GPU返回True，否则返回False
    """
    check_cuda_environment()  # 首先检查CUDA环境
    try:
        logging.info(f"FAISS 库路径: {faiss.__file__}")
        logging.info(f"FAISS 编译信息: {faiss.get_compile_options()}")
        gpu_count = faiss.get_num_gpus()
        logging.info(f"FAISS 报告的 GPU 数量: {gpu_count}")
        
        if gpu_count > 0:
            # 测试每个GPU
            for i in range(gpu_count):
                logging.info(f"测试 GPU {i}")
                res = faiss.StandardGpuResources()  # 尝试创建GPU资源
                logging.info(f"GPU {i} 测试成功")
            logging.info(f"FAISS 检测到 {gpu_count} 个 GPU 设备")
            return True
        else:
            logging.warning("FAISS 未检测到 GPU 设备，将使用 CPU 模式")
            return False
    except AttributeError as e:
        logging.error(f"FAISS GPU 支持错误: {str(e)}")
        logging.error("FAISS 不支持 GPU。请确保安装了支持 GPU 的 FAISS 版本")
        return False
    except Exception as e:
        logging.error(f"检查 GPU 支持时发生未知错误: {str(e)}")
        return False

def main():
    """
    主函数，包含程序的主要逻辑。
    """
    args = get_args()  # 获取命令行参数

    # 记录Python和FAISS版本信息
    logging.info(f"Python 版本: {sys.version}")
    logging.info(f"FAISS 版本: {faiss.__version__}")

    # 如果选择GPU模式，检查GPU支持
    if args.mode == 'gpu':
        if not check_gpu_support():
            logging.error("由于 GPU 不可用，程序将退出。请使用 -m cpu 选项重新运行程序")
            return

    # 尝试创建GPU资源（如果使用GPU模式）
    try:
        res = faiss.StandardGpuResources() if args.mode == 'gpu' else None
        logging.info("成功创建 GPU 资源" if args.mode == 'gpu' else "使用 CPU 模式")
    except AttributeError:
        logging.error("创建 GPU 资源失败。可能的原因：")
        logging.error("1. 安装的 FAISS 版本不支持 GPU")
        logging.error("2. CUDA 环境未正确配置")
        logging.error("请检查 FAISS 安装和 CUDA 配置，或使用 CPU 模式运行")
        return
    except Exception as e:
        logging.error(f"创建 GPU 资源时发生未知错误: {str(e)}")
        return

    # 记录选择的模式、特征提取方法和使用的库
    logging.info(f"选择模式是 {args.mode.upper()}")
    logging.info(f"选择特征提取方法是 {args.feature.upper()}")
    logging.info(f"选择使用的库是 {args.library.upper()}")

    # 加载和预处理数据
    try:
        X, y = createXY(train_folder="../data/train", dest_folder=".", method=args.feature)
        X = np.array(X).astype('float32')
        faiss.normalize_L2(X)  # L2归一化，使向量长度为1
        y = np.array(y)
        logging.info("数据加载和预处理完成。")
        logging.info(f"X 形状: {X.shape}, y 形状: {y.shape}")
    except Exception as e:
        logging.error(f"数据加载或预处理失败：{str(e)}")
        logging.error("请检查数据文件路径是否正确，以及数据格式是否符合要求")
        return

    # 划分数据集为训练集和测试集
    try:
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=2023)
        logging.info("数据集划分为训练集和测试集。")
        logging.info(f"训练集形状: X_train {X_train.shape}, y_train {y_train.shape}")
        logging.info(f"测试集形状: X_test {X_test.shape}, y_test {y_test.shape}")
    except Exception as e:
        logging.error(f"数据集划分失败：{str(e)}")
        return

    # 初始化变量用于存储最佳k值和最高准确率
    best_k = -1
    best_accuracy = 0.0
    k_values = range(1, 6)  # 尝试的k值范围

    # 根据选择的库设置KNN分类器类
    KNNClass = FaissKNeighbors if args.library == 'faiss' else KNeighborsClassifier
    logging.info(f"使用的库为: {args.library.upper()}")

    # 尝试不同的k值，找出最佳的k值
    try:
        for k in tqdm(k_values, desc='寻找最佳k值'):
            # 创建并训练KNN模型
            knn = KNNClass(k=k, res=res) if args.library == 'faiss' else KNNClass(n_neighbors=k)
            knn.fit(X_train, y_train)
            # 在测试集上评估模型
            accuracy = knn.score(X_test, y_test)

            # 更新最佳k值和最高准确率
            if accuracy > best_accuracy:
                best_k = k
                best_accuracy = accuracy
    except Exception as e:
        logging.error(f"模型训练或评估过程中出错：{str(e)}")
        if args.library == 'faiss' and args.mode == 'gpu':
            logging.error("如果使用 FAISS 和 GPU 模式，请确保 FAISS 正确安装并支持 GPU")
        return

    # 输出最佳结果
    logging.info(f'最佳k值: {best_k}, 最高准确率: {best_accuracy}')

# 如果这个脚本是作为主程序运行（而不是被导入），则执行main函数
if __name__ == '__main__':
    main()